# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import logging

from collections import namedtuple

from coin.exchange.shared.kr_rest.currency import SharedCurrency
from coin.proto.coin_query_pb2 import AccountBalance, AccountPosition, ProductPosition, CurrencyBalance

PnlResult = namedtuple('PnlResult', ['reserve', 'quote', 'balance', 'position', 'pnl_balance'])
Balances = namedtuple('Balances', ['reserve', 'balance'])


class PnlCalculator(object):
  def __init__(self, *, scale=1., og, quote_currency, reserves, balance_printer):
    self._base_currencies = reserves
    self._quote_currency = quote_currency
    self._quote_value = None
    self._scale = scale

    self._og = og
    self._balance_printer = balance_printer

    self._reserves = {currency: reserve * self._scale for currency, reserve in reserves.items()}
    self._prices = {}
    self._logger = logging.getLogger(__name__)

  def set_price(self, currency, price):
    self._prices[currency] = price

  def on_book(self, product, book):
    currency = SharedCurrency.FromCurrency(product.base)
    mid_price = 0.5 * (book.ask0().price + book.bid0().price)
    self.set_price(currency.currency, mid_price)

  def set_quote_value(self, value):
    self._quote_value = value

  def get_snapshot(self):
    if self._quote_value is None:
      return
    base_dict = {}
    for base in self._base_currencies:
      reserve = self._reserves.get(base)
      price = self._prices.get(base)
      balance = self._og.get_balance(SharedCurrency.FromStrCurrency(base))
      if reserve is None or price is None or balance is None:
        return
      base_dict['base'] = Balances(reserve=reserve, price=price, balance=balance)
    return base_dict

  def get_snapshot_proto(self):
    if self._quote_value is None:
      return
    if self._quote_currency is None:
      return self._get_snapshot_proto_futures()
    else:
      return self._get_snapshot_proto_spot()

  def _get_snapshot_proto_futures(self):
    pnl_balance = self._balance_printer.calculate_pnl_balance()

    assert len(self._base_currencies) == 1, self._base_currencies
    currency = next(iter(self._base_currencies.keys()))

    # UNSAFE!
    product = self._balance_printer._products[0]
    position, unrealized_pnl = self._balance_printer._get_pos(product)

    res = self._balance_printer.calculate_pnl_balance()

    position_proto = AccountPosition(each_position=[
        ProductPosition(symbol=product.symbol, net_position=position,
                        unrealized_pnl=unrealized_pnl),
    ])
    quote_proto = CurrencyBalance(currency=currency,
                                  quote_price=self._prices[currency],
                                  total=self._og.get_balance(
                                      SharedCurrency.FromStrCurrency(currency)))
    pnl_balance = quote_proto.total
    pnl_balance_proto = CurrencyBalance(currency=currency, total=res[0].total)
    return PnlResult(reserve=None,
                     quote=quote_proto,
                     position=position_proto,
                     balance=None,
                     pnl_balance=pnl_balance_proto)

  def _get_snapshot_proto_spot(self):
    reserve_proto = AccountBalance(each_balance=[
        CurrencyBalance(currency=base, total=self._reserves[base]) for base in self._base_currencies
    ])
    balance = self._og.get_balance(SharedCurrency.FromStrCurrency(self._quote_currency))
    quote_proto = CurrencyBalance(currency=self._quote_currency,
                                  quote_price=self._quote_value,
                                  total=balance)
    pnl_balance = quote_proto.total

    each_balance = []
    for base in self._base_currencies:
      balance = self._og.get_balance(SharedCurrency.FromStrCurrency(base))
      if balance is None:
        self._logger.warn('%s is None', base)
        return
      # Prices here.
      price = self._prices.get(base)
      if price is None:
        self._logger.warn('price (%s) is None', base)
        return
      each_balance.append(CurrencyBalance(currency=base, total=balance, quote_price=price))
    balance_proto = AccountBalance(each_balance=each_balance)
    # Calculate PnlBalance

    for idx, base in enumerate(self._base_currencies):
      proto = balance_proto.each_balance[idx]
      balance = proto.total
      position = balance - self._reserves[base]
      price = proto.quote_price
      pnl_balance += price * position
    pnl_balance_proto = CurrencyBalance(currency=self._quote_currency, total=pnl_balance)
    return PnlResult(reserve=reserve_proto,
                     quote=quote_proto,
                     balance=balance_proto,
                     position=None,
                     pnl_balance=pnl_balance_proto)
